#include <Configuration.h>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <Utilities/Timer.h>
#include <benchmark/p3dfft_bench.h>
using namespace APPNAMESPACE;

template <typename T, unsigned STRIDE>
struct fft_mem
{
  inline static
    void gather(const T* restrict in, T* restrict out, size_t n)
    {
      //while(n)
      //{
      //  *out++=*in; in += STRIDE; --n;
      //}
      //for(int i=n;i>0; --i,in+=STRIDE) *out++=*in;
      for(int i=0; i<n; ++i, in+=STRIDE) *out++=*in;
    }

  inline static
    void gather(const T* restrict in, T* restrict first, T* restrict last)
    {
      //while(first != last)
      for(; first!=last; ++first)
      {
        *first = *in; in += STRIDE;
      }
    }

  inline static
    void scatter(const T* restrict in, T* restrict out, size_t n)
    {
      //while(n)
      //{
      //  *out=*in++; out += STRIDE; --n;
      //}
      //for(int i=n;i>0; --i,out+=STRIDE) *out=*in++;
      for(int i=0;i<n; ++i,out+=STRIDE) *out=*in++;
    }

  inline static
    void scatter(const T* restrict first, const T* restrict last, T* restrict out)
    {
      //while(first != last)
      for(; first!=last;++first)
      {
        *out=*first; out+=STRIDE;
      }
    }

};

int main(int argc, char** argv)
{
  //initialize environment : do not modify 
  mpi::environment env(argc,argv);
  mpi::communicator mycomm;
  OhmmsInfo ohmms("p3dfft",mycomm.rank(),0,1);
  //done with initialization

  app_log() << "Benchmark parallel 3d fft" << std::endl;

  const int default_fft_dim=128;
  const int num_vars=3;
  // default dimension 128 128 128  with 2D distribution
  int nx=default_fft_dim, ny=default_fft_dim, nz=default_fft_dim;
  //number of iterations for timing
  int niters=10;
  // use 0 to use mpi to create 2D node distribution 
  int dims[]={0,0};
  //check correctness
  bool debug=false;

  //check options
  int ic=0;
  while (ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-s")<a.size())
      nx=ny=nz=atoi(argv[++ic]);
    else if(a.find("-mpi_grid")<a.size())
    {
      dims[0]=atoi(argv[++ic]); dims[1]=atoi(argv[++ic]);
    }
    else if(a.find("-i")<a.size())
      niters=atoi(argv[++ic]);
    else if(a.find("-debug")<a.size())
      debug=true;
    else if(a.find("-h") < a.size())
    {
      cout << "For (fft-size)^3 problem on N by M MPI grid " << endl;
      cout << "usgae --s fft-size --mpi_grid N M --i iterations "
           << "[--debug]" << endl;
      mycomm.abort(0);
    }
    ++ic;
  }

  int nproc=mycomm.size();
  if(dims[0]*dims[1]!=nproc)
  {
    app_log() << "#Creating proc. grid with mpi_dims_create\n";
    dims[0]=dims[1]=0;
    MPI_Dims_create(nproc,2,dims);
    if(dims[0] > dims[1]) 
    {
      dims[0] = dims[1];
      dims[1] = nproc/dims[0];
    }
  }

  ///enumeration for the timers
  enum {setup_avg=0,
    forward_avg,
    backward_avg,
    total_avg,
    max_timers
  };
  
  Timer myclock;
  vector<double> timing(max_timers,0.0);

  myclock.restart();
  P3DFFT_benchmark<double> tester(dims,nx,ny,nz);
  timing[setup_avg]=myclock.elapsed();
  mycomm.barrier();

  app_log() << "# P3D FFT dimension " << nx << " " << ny << " " << nz << endl;
  app_log() << "# Using processor grid " << dims[0] << " " << dims[1] << endl;
  if(!mycomm.rank() || (mycomm.rank() == mycomm.size()-1))
  {
    cout << "# Local real-space grid for node= "  << mycomm.rank() << " "
      << tester.isize[0] << " " << tester.isize[1] << " " << tester.isize[2] << endl;
    cout << "# Local k-space grid for node= " << mycomm.rank() << " "
      << tester.fsize[0] << " " << tester.fsize[1] << " " << tester.fsize[2] << endl;
  }

  myclock.restart();
  vector<double> inData(tester.get_local_size_r()*num_vars);
  vector<double> outData(2*tester.get_local_size_k()*num_vars);
  vector<double> outCopy(2*tester.get_local_size_k());
  app_log() << "# Allocation = " << myclock.elapsed()<< " " 
    << inData.size() << " " << outData.size() << endl;

  ///this is the main loop to perform 3 forward 3D FFT and 6 backward 3D FFT
  Timer myclock_main;
  myclock_main.restart();
  size_t r_stride=tester.get_local_size_r();
  size_t k_stride=2*tester.get_local_size_k();

  for(int i=0; i<niters; ++i)
  {
    myclock.restart();
    {
      double* restrict in_ptr=&inData[0];
      for(int k=0; k<num_vars; ++k,in_ptr+=r_stride)
      {
        //fft_mem<3>::gather(in_ptr,&inCopy[0],r_stride);
        //tester.transform_forward(&inCopy[0],&outCopy[0]);
        tester.transform_forward(in_ptr,&outCopy[0]);
        fft_mem<double,3>::scatter(&outCopy[0],&outData[k],k_stride);
      }
    }
    timing[forward_avg]+=myclock.elapsed(); 

    myclock.restart();
    {
      double* restrict in_ptr=&inData[0];
      //double* restrict out_ptr=&outData[0];
      for(int k=0; k<num_vars; ++k, in_ptr+=r_stride)
      {
        fft_mem<double,3>::gather(&outData[k],&outCopy[0],k_stride);
        tester.transform_backward(&outCopy[0],in_ptr);
        //tester.transform_backward(&outCopy[0],&inCopy[0]);
        //fft_mem<3>::scatter(&inCopy[0],in_ptr,r_stride);
      }
    }

    {
      double* restrict in_ptr=&inData[0];
      for(int k=0; k<num_vars; ++k, in_ptr+=r_stride)
      {
        fft_mem<double,3>::gather(&outData[k],&outCopy[0],k_stride);
        tester.transform_backward(&outCopy[0],in_ptr);
        //tester.transform_backward(&outCopy[0],&inCopy[0]);
        //fft_mem<3>::scatter(&inCopy[0],in_ptr,r_stride);
      }
    }
    timing[backward_avg]+=myclock.elapsed(); 
  
    //inData *= factor;
  }
  timing[total_avg]=myclock_main.elapsed();

  //collect timing data
  mpi::reduce(mycomm,timing);

  {//write out for confluence wiki
    const double rk4=4.0;
    double nodefac= rk4/static_cast<double>(mycomm.size());
    timing[setup_avg]*=nodefac;
    timing[forward_avg]*=nodefac/static_cast<double>(niters);
    timing[backward_avg]*=nodefac/static_cast<double>(niters);
    timing[total_avg]*=nodefac/static_cast<double>(niters);

    ostringstream os;
    os << "||p3dfft || fft-size || npx || npy || iterations || setup || forward || backward || total ||\n";
    os << " " << nx << "  " << dims[0] << " " << dims[1] << " " << niters;
    for(int i=0; i<timing.size(); ++i) os << "  " << timing[i];
    app_log() << os.str() << " p3dfft " << endl;
  }

  return 0;
}
